import os
import glob

train_file = ['2_10', '2_11', '2_12', '3_10', '3_11', '3_12', '4_10', '4_11', '4_12', '5_10',
              '5_11', '5_12', '6_7', '6_8', '6_9', '6_10', '6_11', '6_12', '7_7', '7_8', '7_9', '7_11', '7_12']
img_path = './data/Potsdam/3_Ortho_IRRG/'
depth_path = './data/Potsdam/1_DSM_normalisation/'
img_name = glob.glob(img_path+'*IRRG.tif')
train_writer = open('train_test_inputs/potsdam_train.txt','w')
val_writer = open('train_test_inputs/potsdam_val.txt','w')
for i in img_name:
    if len(i[40:-9])==3:
        depth_name = depth_path + 'dsm_potsdam_0' + i[40:-10]+'0'+i[-10] + '_normalized_ownapproach.jpg'
    else:
        depth_name = depth_path+'dsm_potsdam_0'+i[40:-9]+'_normalized_ownapproach.jpg'
    print(depth_name)
    if i[40:-9] in train_file:
        train_writer.write(i+' '+depth_name+'\n')
    else:
        val_writer.write(i + ' ' + depth_name + '\n')
train_writer.close()
val_writer.close()
